import json
from dataclasses import dataclass, field
from typing import Optional, List, Dict, Literal

import uvicorn
import traceback
from fastapi import FastAPI, Body
from langchain_core.messages import AIMessage, HumanMessage
from langchain_ollama import ChatOllama
from pydantic import BaseModel

from src.chat_utils import create_chat
from src.choice_api.choice_api_agent import choice_api_graph
from starlette.responses import StreamingResponse
from src.logger import api_logger

app = FastAPI()


class GraphNode:
    start_entry: Optional[str] = "我需要一些时间思考问题跟现有数据的关系..."
    rewrite_question: Optional[str] = "我需要把问题中的模糊定义标准化,比如时间，"
    chat: Optional[str] = "我需要思考提取一些有效信息来寻找数据"
    tools: Optional[str] = "我需要使用一些工具来回答该问题"
    # api_request: Optional[str] = "获取数据"
    # to_json: Optional[str] = "我需要去访问相应的数据信息..."
    api_request: Optional[str] = "正在获取数据数据中"
    api_parse: Optional[str] = "数据开始清洗"
    data_describe: Optional[str] = "正在进行数据小结"

    @classmethod
    def organization_information(cls, node, value):
        if node == "start_entry":
            return getattr(cls, node)
        if node == "rewrite_question":
            return f"""{getattr(cls, node)}需要把问题修改成:{value.get("rewrite_q", None)}"""
        if node == "chat":
            msg = value.get("api_messages", None)[0]
            if isinstance(msg, AIMessage):
                if len(msg.content) >= 10:
                    return f"""{msg.content}"""
                param_dict = msg.tool_calls[0]["args"] if msg.tool_calls else {}
                parm = ','.join([f"{i}为{j}" for i, j in zip([cls.param_map(key) for key in list(param_dict.keys())], list(param_dict.values()))])
                return f"""{getattr(cls, node)}, {parm}"""
        if node == "tools":
            msg = value.get("api_messages", None)
            if msg:
                return f"""{getattr(cls, node)},查找数据库中的{",".join([cls.map(tool.name) for tool in msg])}信息"""
            return f"""{getattr(cls, node)}"""
        # if node == "to_json":
        #     return f"""{getattr(cls, node)}"""
        if node == "api_request":
            return f"""{getattr(cls, node)}获取数据成功"""
        if node == "api_parse":
            return f"""数据开始清洗"""
        if node == "data_describe":
            print(value)
            clear_results = value.get("clear_results", None)
            if not clear_results:
                clear_results = {
                    'res_total_data': {},
                    'res_total_tip': [],
                    'res_total_table': [],
                    "origin_descriptions": "有什么可以帮助你的吗",
                    "descriptions": "有什么可以帮助你的吗"
                }

            return clear_results

        if node == "common_chat":
            clear_results = value.get("clear_results", None)
            if not clear_results:
                clear_results = {
                    'res_total_data': {},
                    'res_total_tip': [],
                    'res_total_table': [],
                    "origin_descriptions": "有什么可以帮助你的吗",
                    "descriptions": "有什么可以帮助你的吗"
                }
            return clear_results

    @classmethod
    def map(cls, tool_name):
        des = {"get_upland_water_info": "上游来水",
               "get_downstream_water_info": "下游来水",
               "get_lab_internal_info": "实验室在线数据",
               "get_device_check_absence": "查询厂内运行巡检(巡检缺卡/运行缺卡)数据",
               "get_device_check":"查询厂内运行巡检(巡检打卡/运行打卡、巡检正常/巡检异常、巡检异常维修情况)数据",
               "get_electric_consume_record": "查询电量能耗数据",
               "get_internal_online_data": "查询厂内在线数据",
               "get_composite_consume_record": "查询用水量(工业用水、生活用水)能耗数据",
               "get_steam_dosage_record": "查询用蒸汽量能耗数据",
               "get_drug_inventory": "查询药剂管控(出库、入库、库存、供应商)数据",
               "get_sludge_transport_record": "查询污泥外运数据以及污泥含水量数据（按进场/入库时间）",
               # "get_sludge_transport_record_out": "查询污泥外运数据以及污泥外运含水量数据（按出场/出库时间）",
               "get_staff_attendance": "查询人员考勤打卡数据",
               "get_device": "查询设备清单(运行中设备、维修中设备、异常状态设备、停用设备)/设备详情数据",
               }

        return des.get(tool_name, None)

    @classmethod
    def param_map(cls, key):
        param = {
            "startTime": "开始时间",
            "endTime": "结束时间",
            "company_id": "公司名称",
            "target": "指标",
            "sampleId": "检测位置",
            "checkStatus": "状态",
            "shift": "班次",
            "type1": "类型",
            "repairState": "修理状态",
            "abnormal": "厂内是否异常",
            "name": "药品名称",
            "deliverer": "药剂供货商名称",
            "timeType": "进库入库状态",
            "destination": "目的地",
            "timeResult": "打卡情况",
            "state": "设备状态"
        }

        return param.get(key, None)

def handle_post_request(query):
    response = choice_api_graph.stream(
        {"api_messages": [("human", query)]},  # 获取三江化工和传化合成的昨日上游来水和下游来水,查询昨日实验室排放池异常数据
        {"recursion_limit": 10},

    )

    for state in response:
        # print(GraphNode.organization_information(list(state.keys())[0], list(state.values())[0]))
        # print(state)
        # break
        error = list(state.values())[0].get("error", None)
        if error:
            data = {"code": 400, "msg": error, "node": "error"}
            yield json.dumps(data, ensure_ascii=False)
            break
        node = list(state.keys())[0]
        data = {}
        if node in ["start_entry", "rewrite_question", "chat", "tools", "to_json", "api_request", "api_parse"]:
            data = {"code": 200, "msg": GraphNode.organization_information(node, list(state.values())[0]), "node": node}
        if node in ["data_describe", "common_chat"]:
            data = {"code": 200, "data": GraphNode.organization_information(node, list(state.values())[0]), "node": node}
        yield json.dumps(data, ensure_ascii=False) + '\n'


@dataclass
class QueryParam:
    messages: List[Dict] = field(default="", metadata={"description": "用户问题"})


members = ["researcher", "chat"]


# class Router(TypedDict):
#    """Worker to route to next. If no workers needed, route to FINISH."""
#
#    next: Literal[*members]

class Router(BaseModel):
    next: Literal[*members]


def router(query) -> Router:
    # llm = create_chat()
    # system_prompt = (
    #     f"你是一个监督人,当用户发来一个信息,请从{members}的工作者中选出一个符合用户需求的worker,如果是请求或查询数据、清单、情况或台账返回researcher，如果是日常问候返回chat"
    # )
    # messages = [
    #                {"role": "system", "content": system_prompt},
    #            ] + [{"role": "user", "content": query}]
    # response = llm.with_structured_output(Router).invoke(messages)
    # if not response:
    #     return None
    # goto = response.next
    # return goto
    return "researcher"


@app.post("/nl2api")
def nl2api_with_stream(param: QueryParam = Body(...)):
    api_logger.info("POST /nl2api begin")
    messages = param.messages[-1]
    query = messages.get("content", None)
    if router(query) == "researcher":
        try:
            content = handle_post_request(query)
            return StreamingResponse(status_code=200, content=content)
        except Exception as e:
            stack_trace = traceback.format_exc()
            print(stack_trace)
            api_logger.error(f"error: {str(e)}\n{stack_trace}")
            return StreamingResponse(status_code=400, content=json.dumps({"code": 400, "error": str(e)}))
    llm = create_chat()
    response = llm.invoke([HumanMessage(content=query)])
    return StreamingResponse(status_code=200,
                             content=json.dumps({"code": 200, "msg": response.content}, ensure_ascii=False))


if __name__ == '__main__':
    # handle_post_request()
    uvicorn.run(app, host="0.0.0.0", port=55556,log_level="debug")
